from typing import Any, Dict, List, Union

from ..shard.placement_types import (
    DeviceMesh,Replicate,PlacementStrategy,Shard,DTensorSpec,Shard,_Partial,ShardingStrategy
)
import copy
from .strategy_generator import StrategyGenerator


__all__ = ['LayerNormGenerator']


class LayerNormGenerator(StrategyGenerator):
    """
    LayerNormGenerator is a generic class to generate strategies for LayerNorm operation.
    The operation data is defined as `output = input x other + bias`.
    """

    # def validate(self) -> bool:
    #     return super().validate()

    def collate_strategies(self) -> List[ShardingStrategy]:
        '''
        Generate every possible strategies for a LayerNorm node, and record all strategies into the strategies_vector.
        只有三种，复制和shard[2]
        [1,10,768] [768] [768] -> [1,10,768]
        '''
        input_dim = len(self.op_data['input'].data.shape)
        # placements=[[Replicate(),Replicate(),Replicate(),Replicate()],[Shard(input_dim-1),Shard(0),Shard(0),Shard(input_dim-1)],[Shard(input_dim-1),Shard(0),Shard(0),_Partial()]]
        # 将partial替换为replicate
        # placements=[[Replicate(),Replicate(),Replicate(),Replicate()],[Shard(input_dim-1),Shard(0),Shard(0),Shard(input_dim-1)],[Shard(input_dim-1),Shard(0),Shard(0),Replicate()]]
        # 手动删了一种策略 TODO
        placements=[[Replicate(),Replicate(),Replicate(),Replicate()],[Shard(input_dim-1),Shard(0),Shard(0),Replicate()]]
        # 改回有partial
        # placements=[[Replicate(),Replicate(),Replicate(),Replicate()],[Shard(input_dim-1),Shard(0),Shard(0),_Partial()]]
        # 只有复制 TODO
        strategy_list = []
        for placement in placements:
            # 为什么placement要是tuple
            input_spec = DTensorSpec(self.device_mesh, (placement[0],), self.op_data['input'].data)
            weight_spec = DTensorSpec(self.device_mesh, (placement[1],), self.op_data['weight'].data)
            bias_spec = DTensorSpec(self.device_mesh, (placement[2],), self.op_data['bias'].data)
            out_spec = DTensorSpec(self.device_mesh, (placement[3],), self.op_data['output'].data)
            # TODO
            input_specs=[input_spec,weight_spec,bias_spec]
            sharding_specs = PlacementStrategy(input_specs=input_specs,output_specs=out_spec)
            strategy_list.append(ShardingStrategy(name=str(sharding_specs),sharding_specs=sharding_specs,compute_cost=0))
        return strategy_list
